import zipfile
import os
import numpy as np
import cv2
#import matplotlib.pyplot as plt
import json
import random

 #######################################################   
def gaussian2D(shape, sigma=1):
    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m+1,-n:n+1]

    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h

def gaussian_radius(det_size, min_overlap=0.7):
    height, width = det_size

    a1  = 1
    b1  = (height + width)
    c1  = width * height * (1 - min_overlap) / (1 + min_overlap)
    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
    r1  = (b1 + sq1) / 2

    a2  = 4
    b2  = 2 * (height + width)
    c2  = (1 - min_overlap) * width * height
    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
    r2  = (b2 + sq2) / 2

    a3  = 4 * min_overlap
    b3  = -2 * min_overlap * (height + width)
    c3  = (min_overlap - 1) * width * height
    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
    r3  = (b3 + sq3) / 2
    return min(r1, r2, r3)

def draw_umich_gaussian(heatmap, center, radius, k=1):
    diameter = 2 * radius + 1
    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
    
    x, y = int(center[0]), int(center[1])

    height, width = heatmap.shape[0:2]
        
    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    masked_heatmap  = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
    #cv2.imshow('',heatmap)
    #cv2.waitKey(0)
    return heatmap

 #############################################################################   
def Labelencode(img_cv,label,n_classes=10):
    H,W = img_cv.shape[:2]
    hm = np.zeros((128,128,n_classes),dtype=np.float32)
    wh = np.zeros((128,128,2),dtype=np.float32)
    reg = np.zeros((128,128,2),dtype=np.float32)

    x1 = np.asarray(label['left'],dtype=np.float32)
    y1= np.asarray(label['top'],dtype=np.float32)
    w = np.asarray(label['width'],dtype=np.float32)
    h = np.asarray(label['height'],dtype=np.float32)    
    cat = np.asarray(label['label'],dtype=np.int32)    

    for i in range(len(cat)):
        ct_item = ((y1[i]+h[i]/2)/H*128 , (x1[i]+w[i]/2 )/W*128)
        wh_item = (h[i]/H,w[i]/W)  #归一化
        x,y = int(ct_item[1]),int(ct_item[0])
        reg[y,x,:]=[ct_item[0]-y,ct_item[1]-x]
        wh[y,x,:]=[wh_item[0],wh_item[1]]
        radius = gaussian_radius((wh_item[0]*128.,wh_item[1]*128.))
        radius = max(0, int(radius))
        hm_slice = hm[:,:,cat[i]]
        hm_slice = draw_umich_gaussian(hm_slice,(x,y),radius,k=1)
        hm[:,:,cat[i]] = hm_slice
    img_cv = cv2.resize(img_cv,(256,256))
    output = np.concatenate([hm,wh,reg],axis=-1)
    # cv2.imshow('wh_y',cv2.resize(output[...,-4],(400,400)))
    # cv2.imshow('wh_x',cv2.resize(output[...,-3],(400,400)))
    # cv2.imshow('reg_y',cv2.resize(output[...,-2],(400,400)))
    # cv2.imshow('reg_x',cv2.resize(output[...,-1],(400,400)))
    # cv2.waitKey(0)
    return img_cv,output#[hm,wh,reg]


def DataGen(batch_size,img_zip='./data/mchar_train.zip',label_json='./data/mchar_train.json'):
    Z = zipfile.ZipFile(img_zip) #图像压缩包
    image_names = Z.namelist()
    image_names = [_ for _ in image_names if '_MACOSX' not in _ and _[-3:]=='png'] # 提取文件名列表
    file = open(label_json,'r',encoding='utf-8').read()
    label_dict=json.loads(file)
    while True:
        image_names_selected = random.sample(image_names,batch_size)
        img_batch=[]
        label_batch=[]
        for image_name in image_names_selected:
            img_b = Z.read(image_name)
            img_flatten = np.frombuffer(img_b, 'B')
            img_cv = cv2.imdecode(img_flatten, cv2.IMREAD_ANYCOLOR) #获取图像
            #img_plt = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
            label = label_dict[image_name.split('/')[-1]] #获取label信息
            img_decoded,label_decoded=Labelencode(img_cv,label,n_classes=10)
            img_batch.append(img_decoded)
            label_batch.append(label_decoded)
        #num_feat=len(label_batch[0])
        #label_batch = [np.asarray([item[i] for item in label_batch],dtype=np.float32) for i in range(num_feat)] #整形
        yield (np.asarray(img_batch),np.asarray(label_batch))



if __name__ == "__main__":
    for imgs ,labels in DataGen(10):
        pass